/* mem.c : memory manager
 */

#include <kernel.h>
#include <i386.h>

/* Your code goes here */
#define SANITY_BYTE	((uint8) rand())

extern long freemem;
extern char* maxaddr;

struct memHeader {
	uint32 size;
	struct memHeader * prev;
	struct memHeader * next;
	uint8 * sanityCheck;
	uint8 dataStart[0];	
};

/* Function definitions */
void printMemSpace();
void doSanityCheck(struct memHeader * header);
void memoryUnitTest(int numAllocsPerFree, int numFrees, int spacing, int randSizeMultiplier, int useRandSize);

/* 	
*	At kmeminit() time the mem is allocated thusly:
* 	[kernel][freeStart        ][holeMem   ][freeSecondStart                ]
* 	0x0     freemem            HOLESTART   HOLEEND                         maxaddr
*/
struct memHeader *freeStart;
struct memHeader *holeMem;
struct memHeader *freeSecondStart;

/* List pointers updated frequently */
struct memHeader *freeListStart;
struct memHeader *allocListStart;

void kmeminit( void ){
	kprintf( "\nInitializing memory manager.\n");
	if (MEM_UNIT_TEST) kprintf( "freemem: %X, HOLESTART: %X, HOLEEND: %X, maxaddr: %X.\n", 
									freemem, HOLESTART, HOLEEND, maxaddr);
	/*	I can't make a header for the kernel since that would overwrite the first 16 bytes of data. */

	/* 	Write a header for the HOLE.
	*	By casting memory locations to the pointer, we are explicitly defining where the struct data is 
	*	written.
	*	Back up 16 bytes to not overwrite any hole data. 
	*/
	holeMem 		= 	(void*) HOLESTART - sizeof(struct memHeader);	
	holeMem 		-> 	size = HOLEEND - HOLESTART;
	holeMem 		-> 	prev = NULL;
	holeMem 		-> 	next = NULL;
	doSanityCheck(holeMem);
	if (MEM_UNIT_TEST) kprintf( "Hole Defined: Header Addr:  %X, Total Size: %X.\n", 
				holeMem, (holeMem->size + sizeof(struct memHeader)));
	
	/* Start the linked-list of allocated memory regions */
	allocListStart	=	holeMem; 

	/* 	There is a second linked-list for the free memory description
	 	Create a memHeader at freemem. This will describe the first free region. */
	freeStart 				= 	(void*) freemem;
	uint32 freeStartAddr 	= (uint32) freeStart;
	/* Need to allign to paragraph boundaries so truncate 4 LSB, and move to next paragraph if needed */
	if (freeStartAddr & 0x0000000F) {
		freeStart 			= (void*) (freeStartAddr & 0xFFFFFFF0) + 0x10;
		freeStartAddr 		= (uint32) freeStart;
	}
	freeStart 		-> 	size = HOLESTART - freeStartAddr - (2 * sizeof(struct memHeader));
	freeStart 		-> 	prev = NULL;
	freeStart 		-> 	next = (void*) HOLEEND;			/* Point to the next free region after the hole */
	doSanityCheck(freeStart);
	if (MEM_UNIT_TEST) kprintf( "Freespace at end of kernel: Header Addr:  %X, Total Size: %X.\n", 
				freeStart, (freeStart->size + sizeof(struct memHeader)));
		
	/* Start the free memory linked list */
	freeListStart	=	freeStart;
	
	/* And finally the free region at the end. */
	freeSecondStart = 	(void*) HOLEEND;				/* Start writing this header after the hole. */
	freeSecondStart -> 	size = maxaddr - HOLEEND - sizeof(struct memHeader);
	freeSecondStart -> 	prev = (void*) freeStartAddr;
	freeSecondStart -> 	next = NULL;
	doSanityCheck(freeSecondStart);
	if (MEM_UNIT_TEST) kprintf( "Freespace after hole: Header Addr:  %X, Total Size: %X.\n", 
				freeSecondStart, (freeSecondStart->size + sizeof(struct memHeader)));
				
	/* On boot, it is nice to see that the memory is setup correctly */
	printMemSpace();
	
	/* Run the memory system unit test if defined in kernel.h */
	DELAY(MEM_UT_DELAY);
	if (MEM_UNIT_TEST){
		memoryUnitTest( MEM_UT_NUM_ALLOC,  MEM_UT_NUM_FREE,  
						MEM_UT_SPACING, MEM_UT_MIN_SIZE, 
						MEM_UT_DO_RANDSIZE);
		while (TRUE){
			/* Do nothing. */
		}
	}				
}

/* 	Allocate a region of memory.
	If the allocation succeeds, a pointer to the block of memory is 
	returned, otherwise a NULL pointer is returned. 
*/
void *kmalloc( int size ){
	/* Compute the true required amount given that we need space for the header */
	uint16 amnt = (size)/16 + ((size%16)?1:0);
	amnt = amnt*16 + sizeof(struct memHeader);
	if (MEM_UNIT_TEST) kprintf("\nAllocating new memory segment of total size %6X.\n", amnt);	
	
	/* Scan free memory list looking for spot, starting from the header after the kernel space */
	struct memHeader * currentNode = freeListStart;
	while (TRUE){
		if (MEM_UNIT_TEST) kprintf("Currently considering free space of %6X at address %6X.\n", 
										currentNode->size, currentNode);
		if (currentNode->size > (amnt)) {
			/* Here is a free node that is suitable. Let's write the new node here. */
			break;
		}
		if (currentNode->next == NULL) {
			kprintf("MEMORY ERROR! No suitable free space to allocate segment!\n");
			return NULL; /* Return NULL pointer representing no suitable spots. */
		}
		currentNode = currentNode->next; // Increment indexed node */
	}
	
	if (currentNode == NULL){
		/* Something very bad happened */
		kprintf("MEMORY ERROR! Chose NULL for header address!\n");
		return NULL; 
	}
		
	/* 	Now currentNode represents a header location in a suitably sized free region.
	 	Save its data temporarily, so we can update the free-space linked list later */
	struct memHeader * oldFreeNext = currentNode->next;
	struct memHeader * oldFreePrev = currentNode->prev;
	uint32 oldFreeSize = currentNode->size; 
	
	/* 	The link pointers need to now point to ALLOCATED regions, not free regions.	
	 	Start from the holeMem node, since that is a constant point in the allocated list,
	 	and find which allocated nodes should be the next and prev nodes for currentNode. */
	struct memHeader * indexNode = allocListStart;
	/* 	Now we are at the first node. Go forwards towards the end, checking if the node address 
	 	is greater than the new node's address. If it is, then it is the new node's 'next'. */
	while (indexNode < currentNode && indexNode != NULL){
		indexNode = indexNode->next; 
	}
	/* Assign currentNode's pointer to the indexed node */
	currentNode->next = indexNode;
	
	/* Now index to the end of the list to find 'prev'. */
	indexNode = allocListStart;
	while (indexNode->next != NULL){
		indexNode = indexNode->next;
	}
	/* 	Now we have indexed the final node. Index backwards looking for header addresses lower then 
	 	the one for currentNode. This would be the 'prev'. */
	while (indexNode > currentNode && indexNode != NULL){
		indexNode = indexNode->prev; 
	}
	/* Assign */
	currentNode->prev = indexNode;
	
	
	/* Update the adjacent nodes to point to this one */
	if (currentNode->next != NULL) currentNode->next->prev = currentNode;
	if (currentNode->prev != NULL) currentNode->prev->next = currentNode;


	/* 	Adjust the free memory list, in essence deleting the current node from the free list.
	 	Create a new free node at the space after currentNode's data, if there is a free gap. */
	
	if ((currentNode->size + sizeof(struct memHeader)) > amnt){
		/* 	There is a gap of free space suitably large to place a new free-space memHeader.
		 	(Keep in mind that currentNode->size used above still refers to the size 
			of the free region being written to) */
		struct memHeader * proceedingFreespace 	= (void*) (currentNode->dataStart + amnt - sizeof(struct memHeader));	// Free-node address
		/* 	Do paragraph alignment. In hex notation, the LSD needs to be 0.
		 	Mask with the 4 LSB to see if any of them are nonzero. */
		uint32 proceedingFreespaceAddr = (uint32) proceedingFreespace;
		if (proceedingFreespaceAddr & 0x0000000F) {
			if (MEM_UNIT_TEST) kprintf("adjusting paragraph..");
			/* Truncate 4 LSB, and move to next para */
			proceedingFreespace = (void*) (proceedingFreespaceAddr & 0xFFFFFFF0) + 0x10; 
		}
		proceedingFreespace -> prev 			= oldFreePrev;
		proceedingFreespace -> next 			= oldFreeNext;
		proceedingFreespace -> size 			= oldFreeSize - amnt;	// free-space not counting free-node header */
		doSanityCheck(proceedingFreespace);
		
		/* Update the free-node at lower address to point to the new proceeding free-node */
		if (proceedingFreespace -> prev != NULL) {
			proceedingFreespace -> prev -> next = proceedingFreespace;
		}
		if (proceedingFreespace -> next != NULL) {
			proceedingFreespace -> next -> prev = proceedingFreespace;	
		} 
		
		/* If this alloc node used to be the start of freelist, update freelist pointer */
		if (currentNode == freeListStart){
			freeListStart 	=	proceedingFreespace;
		}
	}
	else
	{
		/* 	There is not a large enough gap between the end of this segment and the header for the next.
		 	The following and preceeding freespace headers should be updated to 'delete' the free node being written */
		if (oldFreeNext != NULL) oldFreeNext -> prev = oldFreePrev;
		if (oldFreePrev != NULL) oldFreePrev -> next = oldFreeNext;
		/* If this alloc node used to be the start of freelist, update freelist pointer */
		if (currentNode == freeListStart){
			freeListStart 	=	oldFreeNext;
		}
	}
	

	/* If this is the lowest address alloc node, update alloc list pointer */
	if (currentNode < allocListStart){
		allocListStart	=	currentNode;
	}
	
		
	/* Fill in the rest of the memHeader fields for this allocation. */
	currentNode->size 			= amnt - sizeof(struct memHeader);	/* 'size' field represents the data segment. */
	doSanityCheck(currentNode);
	

	/* Return a pointer to the data */
	if (MEM_UNIT_TEST) kprintf("..success!\n", size);
	return currentNode->dataStart;
}

/* Take a pointer to a previously allocated chunk and free it. */
void kfree( void *ptr ){
	/* Determine the address of the header for the region that is pointed to by the parameter */
	struct memHeader * currentNode = ptr - sizeof(struct memHeader);
	if (MEM_UNIT_TEST) kprintf("Freeing memory segment at header address  %6X. \n", currentNode);
	/* Check if this is sane */
	if (currentNode->sanityCheck != currentNode->dataStart) {
		kprintf("MEM ERROR: Requested free for data at %X but sanity check failed!\n", ptr);
		return;
	}

	/* Update linked lists to delete this node from alloc list and add it to free list, in the right places */
	/* 	Save its data temporarily, so we can update the alloc linked list later */
	struct memHeader * oldAllocNext = currentNode->next;
	struct memHeader * oldAllocPrev = currentNode->prev;
	uint32 oldAllocSize = currentNode->size;
	/* Add this back to the free list */
	struct memHeader * indexNode = freeListStart;
	/* 	Now we are at the first node. Go forwards towards the end, checking if the index node address 
	 	is greater than the freeing node's address. If it is, then it is the new node's 'next'. */
	while (indexNode < currentNode && indexNode != NULL){
		indexNode = indexNode->next; 
	}
	/* Assign currentNode's next pointer to the indexed node */
	currentNode->next = indexNode;
	
	/* Now index to the end of the list to find 'prev'. */
	indexNode = freeListStart;
	while (indexNode->next != NULL){
		indexNode = indexNode->next;
	}
	/* 	Now we have indexed the final node. Index backwards looking for header addresses lower then 
	 	the one for currentNode. This would be the 'prev'. */
	while (indexNode > currentNode && indexNode != NULL){
		indexNode = indexNode->prev; 
	}
	/* Assign */
	currentNode->prev = indexNode;
	/* Update the adjacent nodes to point to this one */
	if (currentNode->next != NULL) currentNode->next->prev = currentNode;
	if (currentNode->prev != NULL) currentNode->prev->next = currentNode;
	/* 	Adjust the allocated memory list, in essence deleting the current node from the alloc list. */	
	if (oldAllocNext != NULL) oldAllocNext -> prev = oldAllocPrev;
	if (oldAllocPrev != NULL) oldAllocPrev -> next = oldAllocNext;
	/* If this node used to be the start of allocList, update allocList pointer */
	if (currentNode == allocListStart){
		allocListStart 	=	oldAllocNext;
	}
	/* If this is the lowest address alloc node, update free list pointer */
	if (currentNode < freeListStart){
		freeListStart	=	currentNode;
	}
	
	/* Coalesce free memory */
	/* 	Scan all free nodes in reverse. If the end of any node is adjacent to the header address of the next node, 
		then erase the next node and take its pointers. Combine the regions into the currentNode */
		
	/* Start by looking at the last node */
	indexNode = freeListStart;
	while (indexNode->next != NULL){
		indexNode = indexNode->next;
	}
	while (indexNode != NULL){
		if ( (struct memHeader*)(indexNode->dataStart + indexNode->size) == indexNode->next){
			/* Adjacent free nodes. Erase the next node and incorporate the space into this node. */
			indexNode->size 		+= 	indexNode->next->size + sizeof(struct memHeader);
			indexNode->next 		= 	indexNode->next->next;
			indexNode->next->prev 	= 	indexNode;
		}
		indexNode = indexNode->prev;
	}
}

void doSanityCheck(struct memHeader * header){
	/* Can update this later to include cool stuff like checksums */
	header -> sanityCheck 	= (uint8*) header->dataStart; 
}

/* ==========================================  UNIT TESTING FUNCTIONS ====================================== */

void printMemHeaderTableLegend(){
	kprintf("   [    addr ][  dataSt ][    prev ][    next ][ dtaSize ][     end ]\n");
}

void printMemHeaderHex(struct memHeader * node){
	kprintf("   [ %7X ][ %7X ][ %7X ][ %7X ][ %7X ][ %7X ]", node, node->dataStart,
	 		node->prev, node->next, node->size, (node->dataStart + node->size));
	if ((uint32)node == (HOLESTART - sizeof(struct memHeader)))	kprintf(" :H");
	if (node == allocListStart)   								kprintf(" :ALS");
	if (node == freeListStart)									kprintf(" :FLS");
	kprintf("\n");
}


void printMemSpace(){
	/* A function to print the contents of the alloc and free lists. */
	kprintf("\nMemory Map:\n");
	printMemHeaderTableLegend();
	
	struct memHeader * indexNode = allocListStart;
	kprintf("Allocated Memory\n");
	/* Go forwards towards the end, printing all allocated nodes. */
	while (indexNode != NULL){
		printMemHeaderHex(indexNode);
		
		indexNode = indexNode->next;
	} 
	
	kprintf("Free Memory\n");
	indexNode = freeListStart;
	/* Go forwards towards the end, printing all free nodes. */
	while (indexNode != NULL){ 
		printMemHeaderHex(indexNode);
		
		indexNode = indexNode->next;
	}
}

void memoryUnitTest(int numAllocsPerFree, int numFrees, int spacing, int randSizeMultiplier, int useRandSize){
	/* This function runs through calling the mem.c functions and displaying the lists each time. */
	kprintf("\n ===== START OF MEMORY UNIT TEST ===== \n");
	
	int i=0;
	int j=0;
	/* Keep pointers to each region so we can free them later */
	uint8 * allocatedAddrs[numAllocsPerFree]; 

	kprintf("Allocating %d segments.\n", numAllocsPerFree);
	for (j; j < numAllocsPerFree; j++ ){
		allocatedAddrs[j] = kmalloc((rand() % 10 * useRandSize + 1) * randSizeMultiplier); 
	}
	
	printMemSpace();
	DELAY(MEM_UT_DELAY);

	kprintf("Freeing %d segments.\n", numFrees);	
	for (i; i < numFrees; i++ ){
		if ((i * spacing) > numAllocsPerFree){
			kprintf("Exceeded address array bounds. Breaking out.\n");
			break;
		}
		if (allocatedAddrs[i * spacing] != NULL){
			kfree(allocatedAddrs[i * spacing]);
		}
	}
	
	printMemSpace();
	DELAY(MEM_UT_DELAY);
	
	kprintf("\n ===== END OF MEMORY UNIT TEST ===== \n");
}


























